from typing import Any

from langchain.retrievers import MilvusRetriever


class CustomMilvusRetriever(MilvusRetriever):
    def get_relevant_documents(self, query, score_threshold=0.5, **kwargs):
        documents = super().get_relevant_documents(query, **kwargs)
        res = [doc for doc in documents if doc.score >= score_threshold]
        return res

